// #include<buddy_system_pmm.h>

// #include<stdio.h>
// #include <list.h>
// #include<assert.h>
// #include <string.h>




// free_area_t free_area;

// #define free_list (free_area.free_list)
// #define nr_free (free_area.nr_free)

// static size_t buddy_physical_size;
// static size_t buddy_virtual_size;
// static size_t buddy_segment_size;
// static size_t buddy_alloc_size;
// static size_t *buddy_segment;
// static struct Page *buddy_physical;
// static struct Page *buddy_alloc;


// static unsigned fixsize(unsigned size) {
//   size |= size >> 1;
//   size |= size >> 2;
//   size |= size >> 4;
//   size |= size >> 8;
//   size |= size >> 16;
//   return size;
// }


// #define REMAINDER(a) ((a)&(fixsize(a)>>1))


// #define BUDDY_ROOT              (1)
// #define BUDDY_LEFT(a)           ((a)<<1)
// #define BUDDY_RIGHT(a)          (((a)<<1)+1)
// #define BUDDY_PARENT(a)         ((a)>>1)
// #define BUDDY_LENGTH(a)         (buddy_virtual_size/ROUNDDOWN(a,2))
// #define BUDDY_BEGIN(a)          (REMAINDER(a)*BUDDY_LENGTH(a))
// #define BUDDY_END(a)            ((REMAINDER(a)+1)*BUDDY_LENGTH(a))
// #define BUDDY_BLOCK(a,b)        (buddy_virtual_size/((b)-(a))+(a)/((b)-(a)))
// #define BUDDY_EMPTY(a)          (buddy_segment[(a)] == BUDDY_LENGTH(a))



// #define MAX(a,b)                ((a)>(b)?(a):(b))
// #define MIN(a,b)                ((a)<(b)?(a):(b))
// #define IS_POWER_OF_TWO(x) (!((x)&((x)-1)))



// static void
// buddy_init_size(size_t n) {
//     assert(n > 1);
//     buddy_physical_size = n;
//     if (n < 512) {
//         buddy_virtual_size = ROUNDUP(n-1,2);
//         buddy_segment_size = 1;
//     } else {
//         buddy_virtual_size = ROUNDDOWN(n,2);
//         buddy_segment_size = buddy_virtual_size*sizeof(size_t)*2/PGSIZE;
//         if (n > buddy_virtual_size + (buddy_segment_size<<1)) {
//             buddy_virtual_size <<= 1;
//             buddy_segment_size <<= 1;
//         }
//     }
//     buddy_alloc_size = MIN(buddy_virtual_size, buddy_physical_size-buddy_segment_size);
// }

// static void
// buddy_init_segment(struct Page *base) {
//     // Init address
//     buddy_physical = base;
//     buddy_segment = KADDR(page2pa(base));
//     buddy_alloc = base + buddy_segment_size;
//     memset(buddy_segment, 0, buddy_segment_size*PGSIZE);
//     // Init segment
//     nr_free += buddy_alloc_size;
//     size_t block = BUDDY_ROOT;
//     size_t alloc_size = buddy_alloc_size;
//     size_t virtual_size = buddy_virtual_size;
//     buddy_segment[block] = alloc_size;
//     while (alloc_size > 0 && alloc_size < virtual_size) {
//         virtual_size >>= 1;
//         if (alloc_size > virtual_size) {
//             // Add left to free list
//             struct Page *page = &buddy_alloc[BUDDY_BEGIN(block)];
//             page->property = virtual_size;
//             list_add(&(free_list), &(page->page_link));
//             buddy_segment[BUDDY_LEFT(block)] = virtual_size;
//             // Switch ro right
//             alloc_size -= virtual_size;
//             buddy_segment[BUDDY_RIGHT(block)] = alloc_size;
//             block = BUDDY_RIGHT(block);
//         } else {
//             // Switch to left
//             buddy_segment[BUDDY_LEFT(block)] = alloc_size;
//             buddy_segment[BUDDY_RIGHT(block)] = 0;
//             block = BUDDY_LEFT(block);
//         }
//     }
//     if (alloc_size > 0) {
//         struct Page *page = &buddy_alloc[BUDDY_BEGIN(block)];
//         page->property = alloc_size;
//         list_add(&(free_list), &(page->page_link));
//     }
// }

// static void
// buddy_init(void) {
//     list_init(&free_list);
//     nr_free = 0;
// }

// static void
// buddy_init_memmap(struct Page *base, size_t n) {
//     assert(n > 0);
//     // Init pages
//     for (struct Page *p = base; p < base + n; p++) {
//         assert(PageReserved(p));
//         p->flags = p->property = 0;
//     }
//     // Init size
//     buddy_init_size(n);
//     // Init segment
//     buddy_init_segment(base);
// }


// static struct Page *
// buddy_alloc_pages(size_t n) {
//     assert(n > 0);
//     struct Page *page;
//     size_t block = BUDDY_ROOT;
//     size_t length = ROUNDUP(n,2);
//     // Find block
//     while (length <= buddy_segment[block] && length < BUDDY_LENGTH(block)) {
//         size_t left = BUDDY_LEFT(block);
//         size_t right = BUDDY_RIGHT(block);
//         if (BUDDY_EMPTY(block)) {                   // Split
//             size_t begin = BUDDY_BEGIN(block);
//             size_t end = BUDDY_END(block);
//             size_t mid = (begin+end)>>1;
//             list_del(&(buddy_alloc[begin].page_link));
//             buddy_alloc[begin].property >>= 1;
//             buddy_alloc[mid].property = buddy_alloc[begin].property;
//             buddy_segment[left] = buddy_segment[block]>>1;
//             buddy_segment[right] = buddy_segment[block]>>1;
//             list_add(&free_list, &(buddy_alloc[begin].page_link));
//             list_add(&free_list, &(buddy_alloc[mid].page_link));
//             block = left;
//         } else if (length & buddy_segment[left]) {  // Find in left (optimize)
//             block = left;
//         } else if (length & buddy_segment[right]) { // Find in right (optimize)
//             block = right;
//         } else if (length <= buddy_segment[left]) { // Find in left
//             block = left;
//         } else if (length <= buddy_segment[right]) {// Find in right
//             block = right;
//         } else {                                    // Shouldn't be here
//             assert(0);
//         }
//     }
//     // Allocate
//     if (length > buddy_segment[block])
//         return NULL;
//     page = &(buddy_alloc[BUDDY_BEGIN(block)]);
//     list_del(&(page->page_link));
//     buddy_segment[block] = 0;
//     nr_free -= length;
//     // Update buddy segment
//     while (block != BUDDY_ROOT) {
//         block = BUDDY_PARENT(block);
//         buddy_segment[block] = buddy_segment[BUDDY_LEFT(block)] | buddy_segment[BUDDY_RIGHT(block)];
//     }
//     return page;
// }




// static void
// buddy_free_pages(struct Page *base, size_t n) {
//     assert(n > 0);
//     struct Page *p = base;
//     size_t length = ROUNDUP(n,2);
//     // Find buddy id 
//     size_t begin = (base-buddy_alloc);
//     size_t end = begin + length;
//     size_t block = BUDDY_BLOCK(begin, end);
//     // Release block
//     for (; p != base + n; p ++) {
//         assert(!PageReserved(p));
//         p->flags = 0;
//         set_page_ref(p, 0);
//     }
//     base->property = length;
//     list_add(&(free_list), &(base->page_link));
//     nr_free += length;
//     buddy_segment[block] = length;
//     // Upadte & merge
//     while (block != BUDDY_ROOT) {
//         block = BUDDY_PARENT(block);
//         size_t left = BUDDY_LEFT(block);
//         size_t right = BUDDY_RIGHT(block);
//         if (BUDDY_EMPTY(left) && BUDDY_EMPTY(right)) {  // Merge
//             size_t lbegin = BUDDY_BEGIN(left);
//             size_t rbegin = BUDDY_BEGIN(right);
//             list_del(&(buddy_alloc[lbegin].page_link));
//             list_del(&(buddy_alloc[rbegin].page_link));
//             buddy_segment[block] = buddy_segment[left]<<1;
//             buddy_alloc[lbegin].property = buddy_segment[left]<<1;
//             list_add(&(free_list), &(buddy_alloc[lbegin].page_link));
//         } else {                                        // Update
//             buddy_segment[block] = buddy_segment[BUDDY_LEFT(block)] | buddy_segment[BUDDY_RIGHT(block)];
//         }
//     }
// }



// static size_t
// buddy_nr_free_pages(void) {
//     return nr_free;
// }


// static void
// macro_check(void) {

//     // Block operate check
//     assert(BUDDY_ROOT == 1);
//     assert(BUDDY_LEFT(3) == 6);
//     assert(BUDDY_RIGHT(3) == 7);
//     assert(BUDDY_PARENT(6) == 3);
//     assert(BUDDY_PARENT(7) == 3);
//     size_t buddy_virtual_size_store = buddy_virtual_size;
//     size_t buddy_segment_root_store = buddy_segment[BUDDY_ROOT];
//     buddy_virtual_size = 16;
//     buddy_segment[BUDDY_ROOT] = 16;
//     assert(BUDDY_LENGTH(6) == 4);
//     assert(BUDDY_BEGIN(6) == 8);
//     assert(BUDDY_END(6) == 12);
//     assert(BUDDY_BLOCK(8, 12) == 6);
//     assert(BUDDY_EMPTY(BUDDY_ROOT));
//     buddy_virtual_size = buddy_virtual_size_store;
//     buddy_segment[BUDDY_ROOT] = buddy_segment_root_store;

//     // Bitwise operate check
//     assert(((0x4000) & (fixsize(0x4000)>>1)) == 0x7FFF);
//     assert(REMAINDER(0x4321) == 0x321);
//     assert(ROUNDUP(0x2321,2) == 0x4000);
//     assert(ROUNDUP(0x2000,2) == 0x2000);
//     assert(ROUNDDOWN(0x4321,2) == 0x4000);
//     assert(ROUNDDOWN(0x4000,2) == 0x4000);

// }

// static void
// size_check(void) {

//     size_t buddy_physical_size_store = buddy_physical_size;
//     buddy_init_size(200);
//     assert(buddy_virtual_size == 256);
//     buddy_init_size(1024);
//     assert(buddy_virtual_size == 1024);
//     buddy_init_size(1026);
//     assert(buddy_virtual_size == 1024);
//     buddy_init_size(1028);    
//     assert(buddy_virtual_size == 1024);
//     buddy_init_size(1030);    
//     assert(buddy_virtual_size == 2048);
//     buddy_init_size(buddy_physical_size_store);   

// }

// static void
// segment_check(void) {

//     // Check buddy segment
//     size_t total = 0, count = 0;
//     for (size_t block = BUDDY_ROOT; block < (buddy_virtual_size<<1); block++)
//         if (BUDDY_EMPTY(block))
//             total += BUDDY_LENGTH(block);
//         else if (block < buddy_virtual_size)
//             assert(buddy_segment[block] == (buddy_segment[BUDDY_LEFT(block)] | buddy_segment[BUDDY_RIGHT(block)]));
//     assert(total == nr_free_pages());

//     // Check free list 
//     total = 0, count = 0;
//     list_entry_t *le = &free_list;
//     while ((le = list_next(le)) != &free_list) {
//         struct Page *p = le2page(le, page_link);
//         count ++, total += p->property;
//     }
//     assert(total == nr_free_pages());

// }

// static void
// alloc_check(void) {

//     // Build buddy system for test
//     size_t buddy_physical_size_store = buddy_physical_size;
//     for (struct Page *p = buddy_physical; p < buddy_physical + 1026; p++)
//         SetPageReserved(p);
//     buddy_init();
//     buddy_init_memmap(buddy_physical, 1026);

//     // Check allocation
//     struct Page *p0, *p1, *p2, *p3;
//     p0 = p1 = p2 = NULL;
//     assert((p0 = alloc_page()) != NULL);
//     assert((p1 = alloc_page()) != NULL);
//     assert((p2 = alloc_page()) != NULL);
//     assert((p3 = alloc_page()) != NULL);

//     assert(p0 + 1 == p1);
//     assert(p1 + 1 == p2);
//     assert(p2 + 1 == p3);
//     assert(page_ref(p0) == 0 && page_ref(p1) == 0 && page_ref(p2) == 0 && page_ref(p3) == 0);

//     assert(page2pa(p0) < npage * PGSIZE);
//     assert(page2pa(p1) < npage * PGSIZE);
//     assert(page2pa(p2) < npage * PGSIZE);
//     assert(page2pa(p3) < npage * PGSIZE);

//     list_entry_t *le = &free_list;
//     while ((le = list_next(le)) != &free_list) {
//         struct Page *p = le2page(le, page_link);
//         assert(buddy_alloc_pages(p->property) != NULL);
//     }

//     assert(alloc_page() == NULL);

//     // Check release
//     free_page(p0);
//     free_page(p1);
//     free_page(p2);
//     assert(nr_free == 3);

//     assert((p1 = alloc_page()) != NULL);
//     assert((p0 = alloc_pages(2)) != NULL);
//     assert(p0 + 2 == p1);

//     assert(alloc_page() == NULL);

//     free_pages(p0, 2);
//     free_page(p1);
//     free_page(p3);

//     struct Page *p;
//     assert((p = alloc_pages(4)) == p0);
//     assert(alloc_page() == NULL);

//     assert(nr_free == 0);

//     // Restore buddy system
//     for (struct Page *p = buddy_physical; p < buddy_physical + buddy_physical_size_store; p++)
//         SetPageReserved(p);
//     buddy_init();
//     buddy_init_memmap(buddy_physical, buddy_physical_size_store);

// }

// static void
// default_check(void) {

//     // Check buddy system
//     macro_check();
//     size_check();
//     segment_check();
//     alloc_check();
    
// }



// const struct pmm_manager buddy_system_pmm_manager = {
//     .name = "buddy_system_pmm_manager",
//     .init = buddy_init,
//     .init_memmap = buddy_init_memmap,
//     .alloc_pages = buddy_alloc_pages,
//     .free_pages = buddy_free_pages,
//     .nr_free_pages = buddy_nr_free_pages,
//     .check = default_check,
// };

#include <pmm.h>
#include <list.h>
#include<stdio.h>
#include <string.h>
#include <buddy_system_pmm.h>

#define LEFT_LEAF(index) ((index) * 2 + 1)
#define RIGHT_LEAF(index) ((index) * 2 + 2)
#define PARENT(index) ( ((index) + 1) / 2 - 1)

#define IS_POWER_OF_2(x) (!((x)&((x)-1)))
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))


free_area_t free_area;

#define free_list (free_area.free_list)
#define nr_free (free_area.nr_free)



struct buddy {
    size_t size; // 子树大小，可管理内存总单元数目
    uintptr_t *longest; // 节点标记，可拓展地址位
    size_t longest_num_page; // 节点分配页表数目的标记
    size_t total_num_page; // 总的分配页表数目
    size_t free_size; // 空闲块的内存单元数目
    struct Page *begin_page; // 管理节点内的指向头页表
};

struct buddy mem_buddy[MAX_NUM_BUDDY];
int num_buddy = 0;

static size_t fix_power_size(size_t size) {
    size |= size >> 1;
    size |= size >> 2;
    size |= size >> 4;
    size |= size >> 8;
    size |= size >> 16;
    return size + 1;
}

static void
buddy_init(void) {
    list_init(&free_list);
    nr_free = 0;
}

static void
buddy_init_memmap(struct Page *base, size_t n) {
    cprintf("n: %d\n", n);
    struct buddy *buddy = &mem_buddy[num_buddy++];

    size_t v_size = fix_power_size(n);
    size_t excess = v_size - n;
    size_t v_alloced_size = fix_power_size(excess);

    buddy->size = v_size;
    buddy->free_size = v_size - v_alloced_size;
    buddy->longest = page2kva(base);
    buddy->begin_page = pa2page(PADDR(ROUNDUP(buddy->longest + 2 * v_size * sizeof(uintptr_t), PGSIZE)));
    buddy->longest_num_page = buddy->begin_page - base;
    buddy->total_num_page = n - buddy->longest_num_page;

    size_t node_size = buddy->size * 2;

    for (int i = 0; i < 2 * buddy->size - 1; i++) {
        if (IS_POWER_OF_2(i + 1)) {
            node_size /= 2;
        }
        buddy->longest[i] = node_size;
    }

    int index = 0;
    while (1) {
        if (buddy->longest[index] == v_alloced_size) {
            buddy->longest[index] = 0;
            break;
        }
        index = RIGHT_LEAF(index);
    }

    while (index) {
        index = PARENT(index);
        buddy->longest[index] = MAX(buddy->longest[LEFT_LEAF(index)], buddy->longest[RIGHT_LEAF(index)]);
    }

    struct Page *p = buddy->begin_page;
    for (; p != base + buddy->free_size; p ++) {
        assert(PageReserved(p));
        p->flags = p->property = 0;
        set_page_ref(p, 0);
    }
}

static struct Page *
buddy_alloc_pages(size_t n) {
    assert(n > 0);
    if (!IS_POWER_OF_2(n))
        n = fix_power_size(n);

    size_t index = 0;
    size_t node_size;
    size_t offset = 0;

    struct buddy *buddy = NULL;
    for (int i = 0; i < num_buddy; i++) {
        if (mem_buddy[i].longest[index] >= n) {
            buddy = &mem_buddy[i];
            break;
        }
    }

    if (!buddy) {
        return NULL;
    }

    for (node_size = buddy->size; node_size != n; node_size /= 2) {
        if (buddy->longest[LEFT_LEAF(index)] >= n)
            index = LEFT_LEAF(index);
        else
            index = RIGHT_LEAF(index);
    }

    buddy->longest[index] = 0;
    offset = (index + 1) * node_size - buddy->size;

    while (index) {
        index = PARENT(index);
        buddy->longest[index] = MAX(buddy->longest[LEFT_LEAF(index)], buddy->longest[RIGHT_LEAF(index)]);
    }

    buddy->free_size -= n;

    return buddy->begin_page + offset;
}


static void
buddy_free_pages(struct Page *base, size_t n) {
    struct buddy *buddy = NULL;

    for (int i = 0; i < num_buddy; i++) {
        struct buddy *t = &mem_buddy[i];
        if (base >= t->begin_page && base < t->begin_page + t->size) {
            buddy = t;
        }
    }

    if (!buddy) return;

    unsigned node_size, index = 0;
    unsigned left_longest, right_longest;
    unsigned offset = base - buddy->begin_page;

    assert(offset >= 0 && offset < buddy->size);

    node_size = 1;
    index = offset + buddy->size - 1;

    for (; buddy->longest[index]; index = PARENT(index)) {
        node_size *= 2;
        if (index == 0)
            return;
    }

    buddy->longest[index] = node_size;
    buddy->free_size += node_size;

    while (index) {
        index = PARENT(index);
        node_size *= 2;

        left_longest = buddy->longest[LEFT_LEAF(index)];
        right_longest = buddy->longest[RIGHT_LEAF(index)];

        if (left_longest + right_longest == node_size)
            buddy->longest[index] = node_size;
        else
            buddy->longest[index] = MAX(left_longest, right_longest);
    }

}


static size_t
buddy_nr_free_pages(void) {
    size_t total_free_pages = 0;
    for (int i = 0; i < num_buddy; i++) {
        total_free_pages += mem_buddy[i].free_size;
    }
    return total_free_pages;
}


static void
buddy_check(void) {
    size_t total = buddy_nr_free_pages();
    cprintf("total: %d\n", total);

    struct Page *p0 = alloc_page();
    assert(p0 != NULL);
    assert(buddy_nr_free_pages() == total - 1);
    assert(p0 == mem_buddy[0].begin_page);

    struct Page *p1 = alloc_page();
    assert(p1 != NULL);
    assert(buddy_nr_free_pages() == total - 2);
    assert(p1 == mem_buddy[0].begin_page + 1);

    assert(p1 == p0 + 1);

    buddy_free_pages(p0, 1);
    buddy_free_pages(p1, 1);
    assert(buddy_nr_free_pages() == total);

    p0 = buddy_alloc_pages(11);
    assert(buddy_nr_free_pages() == total - 16);

    p1 = buddy_alloc_pages(100);
    assert(buddy_nr_free_pages() == total - 144);

    buddy_free_pages(p0, -1);
    buddy_free_pages(p1, -1);
    assert(buddy_nr_free_pages() == total);

    p0 = buddy_alloc_pages(total);
    assert(p0 == NULL);

    // debug_buddy_tree(7, "221, init");
    p0 = buddy_alloc_pages(512);
    // debug_buddy_tree(7, "221, alloc 512");
    assert(buddy_nr_free_pages() == total - 512);

    p1 = buddy_alloc_pages(1024);
    // debug_buddy_tree(7, "225, alloc 1024");
    assert(buddy_nr_free_pages() == total - 512 - 1024);

    struct Page *p2 = buddy_alloc_pages(2048);
    // debug_buddy_tree(7, "229, alloc 2048");
    assert(buddy_nr_free_pages() == total - 512 - 1024 - 2048);

    struct Page *p3 = buddy_alloc_pages(4096);
    // debug_buddy_tree(7, "233, alloc 4096");
    assert(buddy_nr_free_pages() == total - 512 - 1024 - 2048 - 4096);

    struct Page *p4 = buddy_alloc_pages(8192);
    // debug_buddy_tree(7, "237, alloc 8192");
    assert(buddy_nr_free_pages() == total - 512 - 1024 - 2048 - 4096 - 8192);

    struct Page *p5 = buddy_alloc_pages(8192);
    // debug_buddy_tree(7, "241, alloc 8192");
    assert(buddy_nr_free_pages() == total - 512 - 1024 - 2048 - 4096 - 8192 - 8192);

    buddy_free_pages(p0, -1);
    buddy_free_pages(p1, -1);
    buddy_free_pages(p2, -1);
    buddy_free_pages(p3, -1);
    buddy_free_pages(p4, -1);
    buddy_free_pages(p5, -1);

    assert(buddy_nr_free_pages() == total);

}


static void
debug_buddy_tree(int level, char *label) {
    cprintf("\ndebug buddy tree: %s\n", label);

    int num = 1;
    int index = 0;
    for (int i = 0; i < level; i++) {
        int cur_num = 0;
        while (1) {
            cprintf("%d ", mem_buddy[0].longest[index++]);
            cur_num++;
            if (cur_num == num) {
                cprintf("\n");
                cur_num = 0;
                num *= 2;
                break;
            }
        }
    }
    cprintf("debug buddy tree end\n\n");
}


    const struct pmm_manager buddy_system_pmm_manager = {
            .name = "buddy_system_pmm_manager",
            .init = buddy_init,
            .init_memmap = buddy_init_memmap,
            .alloc_pages = buddy_alloc_pages,
            .free_pages = buddy_free_pages,
            .nr_free_pages = buddy_nr_free_pages,
            .check = buddy_check,
    };